#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright 2024 Huawei Technologies Co., Ltd
import numpy as np
import os

rank_dim = 8
rank_m = 512
rank_k = 5120
rank_n = 640

def gen_golden_data_simple():
    if not os.path.exists("input"):
        os.mkdir("input")
    if not os.path.exists("output"):
        os.mkdir("output")

    input_x1 = []
    input_x2 = []
    for i in range(rank_dim):
        x1 = np.random.uniform(-3, 3, [rank_m, rank_k]).astype(np.float16)
        x2 = np.random.uniform(-3, 3, [rank_k, rank_n]).astype(np.float16)
        x1.tofile("./input/input_x1_{}.bin".format(i))
        x2.tofile("./input/input_x2_{}.bin".format(i))
        input_x1.append(x1)
        input_x2.append(x2)

    golden_gather_out = np.concatenate(input_x1, axis=0)
    for i in range(rank_dim):
        golden_gather_out.tofile("./output/golden_gather_out_{}.bin".format(i))
        out = np.matmul(golden_gather_out.astype(np.float32), input_x2[i].astype(np.float32)).astype(np.float16)
        out.tofile("./output/golden_out_{}.bin".format(i))

if __name__ == "__main__":
    gen_golden_data_simple()
